function [x,normk,it,nrm] = solverPI(in)
%% function [x,normk,it,nrm] = solverPI(input)
%   Solves the coarse mesh diffusion problem via the power 
%   iteration technique.
% needs:
%   * function to compute leakage
%   * a way to normalize fission rates
%   * a way to plot outgoing partial currents for a given element
% j. roberts 2/22/2010
if(in.maxit>1)
disp('-------------------------------------------')
disp('--       beginning power iterations      --')
disp('-------------------------------------------')
end

M     = in.M;
order = in.order;
numg  = in.numg;
numel = in.numel;
x     = in.x0;

N   = length(M(1,:));  % should be 8 for 0th, 16 for 1st, etc.
j1  = M*x(1:end-1, 1);
k1  = x(end, 1);

k0 = 0; 
it = 1;
jwarn=0;
jjit=0;
maxjit = 6000;  % probably make a runtime parameter
jold=zeros(N,1); 
%nrmm = norm( resid2D( [j1' k1 1.0]', in) );

% initialize response matrices
R = zeros( numel*numg*(order+1)*4 );
F = zeros( 1, numel*numg*(order+1)*4 );
A = zeros( 1, numel*numg*(order+1)*4 );
L = zeros( 4, numel*numg*(order+1)*4 );
lenR = numg*(order+1)*4;

% build mm
for i = 1:length(M(:,1))
    if ( sum(M(:,i)) + sum(M(i,:)) )==0
        mm(i)=1;
    else
        mm(i)=0;
    end
end
mm = diag(mm(1:numg*(order+1):end));

nrm(1) = norm( resid2D([j1' k1 1.0]',in,j1) );
nrmk(1) = 1;

while nrm(it) > in.tol(1)*nrm(it) + in.tol(2) && it < in.maxit+1
    if (in.maxit>1)
    disp(['  iteration number: ', num2str(it), ' norm(x) = ', num2str(nrm(it))])
    end

    % get response matrices for current k value
    [Rb Fb Ab Lb]  = resp2D(k1,in);
    % this gets the info for each element type; now, need to place
    I = length(in.elements(:,1));
    J = length(in.elements(1,:));
    k = 0;
    % this should go in a function, or at the end of resp2D;
    for j = 1:J      % rows
        for i = 1:I  % cols
            if in.elements(i,j) > 0
                % I am the kth element
                k = k + 1;
                type = in.elements(i,j);
            else
                % else I'm just void
                continue
            end
            R( ((k-1)*lenR+1):(k*lenR),((k-1)*lenR+1):(k*lenR) ) = Rb(type,:,:);
            F(1, ((k-1)*lenR+1):(k*lenR) ) = Fb(type,:);
            A(1, ((k-1)*lenR+1):(k*lenR) ) = Ab(type,:); 
        end
    end
    R = sparse(R);
    j0 = zeros(N,1);
    jit = 0;
    jold = j1;
    MR=M*R;
    while norm(j1-j0) > 1e-12 && jit < maxjit
        sumj   = sqrt(j1'*j1);
        j0      = j1;
        j1      = MR*j1/sumj;     
        jit     = jit + 1;
    end
    if (jit==maxjit)
        jjit=jjit+1;
        jwarn=max(1,jwarn);
    end
    k0 = k1;
    % keep in mind j's are INCOMING currents
    jout = R*j1; 
    sqrtjout = sqrt(jout'*jout);
    jout = jout/sqrtjout;
    gain = F*j1;  
    % need to compute leakage.
    L = 0;
    tmp = numg*(order+1);
    k=0;
    for j = 1:J      % rows
        for i = 1:I  % cols
            if in.elements(i,j) > 0
                % I am the kth element
                k = k + 1;
                type = in.elements(i,j);
            else
                % else I'm just void
                continue
            end
            
            myL(:,:) = Lb(type,:,:);
            my1 = myL(1,:);
            my2 = myL(2,:);
            my3 = myL(3,:);
            my4 = myL(4,:);
            % for an element, we have
            Ltmp = [ my1 my2 my3 my4
                     my2 my1 my4 my3
                     my4 my3 my1 my2
                     my3 my4 my2 my1 ];

            % build the L matrix
            LK( (k-1)*4+1:(k-1)*4+4, (k-1)*4*tmp+1:k*4*tmp ) = Ltmp;
        end
    end
    L = sum( mm*(j1'*LK')' );

    loss = (A*j1+L);% + abs(1-mm(1,1))*LL;% + abs(1-mm(2,2))*LR +...
                %  abs(1-mm(3,3))*LB + abs(1-mm(4,4))*LT;
    % NOTE 3.17.2010 -- because I already renormed jout by sqrt term, I've
    %  got to account for that on the loss term, too, which uses only j-
    lam = max(abs(eig(full(R*M))));
    loss2 = loss/sqrtjout*lam;
    % temp norm F !!!
    FF = [ R*M - lam*eye(length(R(:,1)))    zeros(length(R(:,1)),1)
           F*M                                             -loss2 ];
    z1 = FF*([jout' k1]');
    z2 = 0.5 - 0.5*(jout'*jout);                        
    k1 = gain/loss;
    it = it + 1;
    nrm(it) = norm( [z1;z2] );
    
    %nrm(it) = norm( resid2D([jout' k1 max(eig(full(R*M)))]',in, j1) );
    normk(it) = norm(k1-k0);
end

if(jwarn>0)
    disp('warning: at least one inner did not converge')
    disp(['# of kits: ',num2str(it),' # of jjits ',num2str(jjit)])
end
x = [jout' k1 1.0]';
end

function z = resid2D(x,in,j1)
%% function z = resid2D(x,input)
% This function computes the nonlinear residual
% Here, x(1:end-2) corresponds to the legendre coefficients
%       x(end-1) is keff
%       x(end)   is the current eigenvalue (lambda)
M = in.M;
order = in.order;
numel = in.numel;
numg = in.numg;
%j1=M*x(1:end-2);
%j1 = x(1:end-2);
[R,F,A,L] = combresp2D(in,[j1' x(end-1) x(end)]');
loss = A*M*x(1:end-2)+L;

%compute the nonlinear function, FF(x)*x = 0
FF = [ R*M-x(end)*eye(length(R(:,1)))  zeros(length(R(:,1)),1)
       F*M                            -loss ];
z1 = FF*(x(1:end-1));
z2 = 0.5 - 0.5*x(1:end-2)'*x(1:end-2);
z  = [z1' z2']';
end % resid2D


function [R,F,A,L] = combresp2D(in,x)
M = in.M;
order = in.order;
numel = in.numel;
numg = in.numg;
% initialize response matrices
R = zeros( numel*numg*(order+1)*4 );
F = zeros( 1, numel*numg*(order+1)*4 );
A = zeros( 1, numel*numg*(order+1)*4 );
L = zeros( 4, numel*numg*(order+1)*4 );
lenR = numg*(order+1)*4;
% build mm
for i = 1:length(M(:,1))
    if ( sum(M(:,i)) + sum(M(i,:)) )==0
        mm(i)=1;
    else
        mm(i)=0;
    end
end
mm = diag(mm(1:numg*(order+1):end));
%% Get response matrices
[Rb Fb Ab Lb]  = resp2D(x(end-1),in);
% this gets the info for each element type; now, need to place
I = length(in.elements(:,1));
J = length(in.elements(1,:));
k = 0;
for j = 1:J      % rows
    for i = 1:I  % cols
        if in.elements(i,j) > 0
            % I am the kth element
            k = k + 1;
            type = in.elements(i,j);
        else
            % else I'm just void
            continue
        end
        R( ((k-1)*lenR+1):(k*lenR),((k-1)*lenR+1):(k*lenR) ) = Rb(type,:,:);
        F(1, ((k-1)*lenR+1):(k*lenR) ) = Fb(type,:);
        A(1, ((k-1)*lenR+1):(k*lenR) ) = Ab(type,:);
    end
end
R = sparse(R);
tmp = numg*(order+1);
k=0;
%% Get the leakage
for j = 1:J      % rows
    for i = 1:I  % cols
        if in.elements(i,j) > 0
            % I am the kth element
            k = k + 1;
            type = in.elements(i,j);
        else
            % else I'm just void
            continue
        end
        
        myL(:,:) = Lb(type,:,:);
        my1 = myL(1,:);
        my2 = myL(2,:);
        my3 = myL(3,:);
        my4 = myL(4,:);
        % for an element, we have
        Ltmp = [ my1 my2 my3 my4
            my2 my1 my4 my3
            my4 my3 my1 my2
            my3 my4 my2 my1 ];
        
        % build the L matrix
        LK( (k-1)*4+1:(k-1)*4+4, (k-1)*4*tmp+1:k*4*tmp ) = Ltmp;
    end
end
L = sum( mm*(x(1:end-2)'*LK')' );
end % combresp2D
